import torch
import torch.nn as nn
import torch.nn.init as init
import torch.nn.functional as F

from torch.utils import model_zoo
from torchvision import models
class Mobile_Unet(nn.Module):

    def __init__(self, num_classes,alpha=0.5):
        super(Mobile_Unet, self).__init__()
        def conv_bn(inp, oup, stride):
            return nn.Sequential(
                nn.Conv2d(inp, oup, 3, stride, 1, bias=False),
                nn.BatchNorm2d(oup),
                nn.ReLU(inplace=True)
            )

        def conv_dw(inp, oup, stride):
            return nn.Sequential(
                nn.Conv2d(inp, inp, 3, stride, 1, groups=inp, bias=False),
                nn.BatchNorm2d(inp),
                nn.ReLU(inplace=True),

                nn.Conv2d(inp, oup, 1, 1, 0, bias=False),
                nn.BatchNorm2d(oup),
                nn.ReLU(inplace=True),
            )
        self.b00 = nn.Sequential(
            conv_bn(  3,  int(32*alpha), 2),
        )

        self.b01 = nn.Sequential(
            conv_dw(int(32*alpha),  int(64*alpha), 1),)

        self.b03 = nn.Sequential(
            conv_dw(int(64*alpha), int(128*alpha), 2),
            conv_dw(int(128*alpha), int(128*alpha), 1),)

        self.b05 = nn.Sequential(
            conv_dw(int(128*alpha), int(256*alpha), 2),
            conv_dw(int(256*alpha), int(256*alpha), 1),)

        self.b11 = nn.Sequential(
            conv_dw(int(256*alpha), int(512*alpha), 2),
            conv_dw(int(512*alpha), int(512*alpha), 1),
            conv_dw(int(512*alpha), int(512*alpha), 1),
            conv_dw(int(512*alpha), int(512*alpha), 1),
            conv_dw(int(512*alpha), int(512*alpha), 1),
            conv_dw(int(512*alpha), int(512*alpha), 1),
        )

        self.b13 = nn.Sequential(
            conv_dw(int(512*alpha), int(1024*alpha), 2),
            conv_dw(int(1024*alpha), int(1024*alpha), 1),
        )

        self.ConvTranspose1= nn.ConvTranspose2d(int(1024*alpha), int(512*alpha), 2, stride=2)
        self.b14 = conv_dw(int(1024*alpha), int(512*alpha), 1)

        self.ConvTranspose2= nn.ConvTranspose2d(int(512*alpha), int(256*alpha), 2, stride=2)
        self.b15 = conv_dw(int(512*alpha), int(256*alpha), 1)

        self.ConvTranspose3= nn.ConvTranspose2d(int(256*alpha), int(128*alpha), 2, stride=2)
        self.b16 = conv_dw(int(256*alpha), int(128*alpha), 1)

        self.ConvTranspose4= nn.ConvTranspose2d(int(128*alpha), int(64*alpha), 2, stride=2)
        self.b17 = conv_dw(int(128*alpha), int(64*alpha), 1)

        self.b18 = conv_bn(int(96*alpha), int(32*alpha), 1)

        self.final = nn.Conv2d(int(32*alpha), num_classes, 1)

    def forward(self, x):
        b00 =self.b00(x)
        b01 =self.b01(b00)
        b03 =self.b03(b01)
        b05 =self.b05(b03)
        b11 =self.b11(b05)
        b13 =self.b13(b11)

        up1 = torch.cat([self.ConvTranspose1(b13),b11],1)
        b14 = self.b14(up1)

        up2 = torch.cat([self.ConvTranspose2(b14),b05],1)
        b15 = self.b15(up2)

        up3 = torch.cat([self.ConvTranspose3(b15),b03],1)
        b16 = self.b16(up3)

        up4 = torch.cat([self.ConvTranspose4(b16),b01],1)
        b17 = self.b17(up4)

        up5 = torch.cat([b17,b00],1)
        b18=self.b18(up5)
        return F.upsample_bilinear(self.final(b18),scale_factor=2)
